"""
nnunet/network_architecture/generic_UNet.py
    - Generic_UNet(SegmentationNetwork)
        - nnunet/network_architecture/neural_network.py
        - SegmentationNetwork
    - 这个UNet是融合3D和2D的经典UNet，但和常规的不一致，不要相提并论

由于这个框架的各类调用比较多，八成会比较依赖torch本身，所以导出TorchScript和Onnx
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import torch
from torch import nn
from nnunet.network_architecture.generic_UNet import Generic_UNet
from nnunet.training.model_restore import load_model_and_checkpoint_files

trainer, params = load_model_and_checkpoint_files(
    "../resource/Task670_conus_resized_2000/nnUNetTrainerV2__nnUNetPlansv2.1",
    0,
    mixed_precision=True,
    checkpoint_name="model_best"
)
assert type(trainer.network) == Generic_UNet, "Not Impl for other architecture"
trainer.load_checkpoint_ram(params[0], False)
trainer.network.eval()

# 减少输出
trainer.network._deep_supervision = False
trainer.network.do_ds = False

# 调整warning
# https://github.com/pytorch/pytorch/issues/75252
# 滞后: 将nnUNetTrainV2.py 修改 norm_op_kwargs = {'eps': 1e-5, 'affine': True, "track_running_stats": True}
def print_module_training_status(module):
    if isinstance(module, nn.Conv2d) or isinstance(module, nn.Conv3d) or isinstance(module, nn.Dropout3d) or \
            isinstance(module, nn.Dropout2d) or isinstance(module, nn.Dropout) or isinstance(module, nn.InstanceNorm3d) \
            or isinstance(module, nn.InstanceNorm2d) or isinstance(module, nn.InstanceNorm1d) \
            or isinstance(module, nn.BatchNorm2d) or isinstance(module, nn.BatchNorm3d) or isinstance(module,
                                                                                                      nn.BatchNorm1d):
        print(str(module), module.training)
        m.train(False)

for m in trainer.network.modules():
    print_module_training_status(m)

for param in trainer.network.parameters():
    param.requires_grad = False

# 注意根据实际情况调整
input_tensor = torch.randn([1, 3, 512, 512], requires_grad=False)

# save torch script
with torch.no_grad():
    jit_model = torch.jit.trace(trainer.network, input_tensor)
    jit_model.save('v1_jit.pth')

# save onnx
with torch.no_grad():
    torch.onnx.export(trainer.network,
                      input_tensor,
                      "v1_onnx.onnx",
                      verbose=True,
                      input_names=['input:0'],
                      output_names=['output:0'],
                      opset_version=17)
